"""
@author: j.l.liu
@create: 
@description: 把任意 txt / md / pdf 等文件变成向量存到本地 chroma_db 目录
1、LangChain框架 负责“把问题拆成步骤、找资料、调用外部工具”
2、DeepSeek模型 负责“真正动脑”：理解问题、生成答案或代码、做数学推理
两者通过标准接口（OpenAI-compatible API 或 langchain-deepseek 包）对接
"""
import os

from langchain.document_loaders import DirectoryLoader, TextLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
docs_path = './docs'
# 1) 目录不存在 → 直接提示并退出
if not os.path.isdir(docs_path):
    print(f"❌ 目录不存在：{docs_path}")
    exit(1)
# 2. 加载目录 ./docs 下所有 .txt 文件
loader = DirectoryLoader('./docs', glob="**/*.txt", loader_cls=TextLoader,loader_kwargs={'encoding': 'utf-8'})
docs = loader.load()
# 2) 没有 txt 文件 → 提示并退出
if not docs:
    print(f"⚠️ 目录 {docs_path} 下没有找到任何 .txt 文件")
    exit(0)
print(f"✅ 已找到 {len(docs)} 份文档，准备向量化…")
# 3. 把长文档按 500 字符长度切块，相邻块重叠 50 字符，防止语义断裂
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
texts = text_splitter.split_documents(docs)
# 4. 指定中文句向量模型（下载到 ~/.cache/huggingface）
embeddings = HuggingFaceEmbeddings(
    model_name="./models/text2vec-base-chinese",  # 本地路径
    model_kwargs={"local_files_only": True}       # 不再联网
)

# 5. 将切分后的文本向量化，并存入本地目录 ./chroma_db
#    Chroma 会自动把文本 + 向量 + 元数据写入 SQLite 文件
db = Chroma.from_documents(texts, embeddings, persist_directory="./chroma_db")

# 6. 立即落盘，确保下次启动可以直接加载
db.persist()

# 7. 打印统计信息
print("知识库已生成，共 %d 条文档" % len(texts))